import asyncio
from concurrent.futures import ThreadPoolExecutor
from enum import Enum
from typing import Type

from api.plugins.client import PluginClient
from api.plugins.consts import ToolName, PLUGIN_CONFIG, PLUGIN_PARAMETER, RAG_CONCURRENCY

tools_map = {}
rag_tools = []
plugins = []


def is_string_in_enum(value: str, enum_class: Type[Enum]) -> bool:
    return any(value == item.value for item in enum_class)


def get_plugin_path(path: str) -> str:
    return path.format(**PLUGIN_PARAMETER)


def run_async_task(loop, plugin, tool, *args):
    asyncio.set_event_loop(loop)
    return loop.run_until_complete(plugin.exec(tool.name, *args))


def exec_plugin(tool_name: ToolName, *args):
    # RAG时，多路召回并行处理
    if tool_name == ToolName.RAG:
        futures = []
        with ThreadPoolExecutor(max_workers=RAG_CONCURRENCY) as executor:
            for plugin, loop, tool in rag_tools:
                future = executor.submit(run_async_task, loop, plugin, tool, *args)
                futures.append((future, tool.name))

        results = []
        for future, name in futures:
            result = future.result()
            setattr(result, 'tool_name', name)
            results.append(result)

        return results

    plugin, loop, tool = tools_map[tool_name.value]
    asyncio.set_event_loop(loop)
    return loop.run_until_complete(plugin.exec(tool.name, *args))


def generate_plugin(plugin_path: str):
    plugin = PluginClient(get_plugin_path(plugin_path))
    # 这里为了解决异步问题，需要自管理异步事件
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    plugins.append((plugin, loop))

    loop.run_until_complete(plugin.connect())

    resp = loop.run_until_complete(plugin.session.list_tools())

    for tool in resp.tools:
        tool_name = tool.name

        # 固定tool记录，其它tool定义为RAG
        if is_string_in_enum(tool.name, ToolName):
            tools_map[tool_name] = (plugin, loop, tool)
        else:
            rag_tools.append((plugin, loop, tool))


def generate_plugins():
    for path in PLUGIN_CONFIG:
        generate_plugin(path)


if __name__ == '__main__':
    generate_plugins()

    # print(exec_plugin(ToolName.RELATION_CHECK, {
    #     'question': '智界R7的价格是多少？'.encode('unicode_escape'),
    #     'knowledge': ['目前智界R7的在售车型如下：'.encode('unicode_escape'),
    #                   '目前智界R7的在售车型如下：智界R7'
    #                   '这款车拥有颠覆想象的造型设计、卓越的性能表现及多维的配置选项为一体，它在让每一个热爱生活的用户都能轻松踏上梦想中的旅程'
    #                   '。欢迎各位车友在评论区留言，说说您对智界S7的看法。'.encode('unicode_escape'),
    #                   '智界R7是奇瑞汽车与华为联合打造的一款纯电轿跑SUV，也是智界品牌的第二款车型，共推出四款车型，智界R7 Pro版售价为25.98'
    #                   '万元，智界R7 Max 82度电池版售价为29.98万元，智界R7 Max 100度电池版售价为31.98万元，智界R7 '
    #                   'Ultra版售价为33.98万元。智界R7集颠覆想象的造型设计、卓越的性能表现及多维的配置选项为一体，代表着智界品牌的又一次出发'
    #                   '。下面小编就给大家介绍一下智界R7这款车。'.encode('unicode_escape')
    #                   ]}))
    # print(exec_plugin(ToolName.RAG, {'query': '你好', 'slang': []}))
    # question_classification = exec_plugin(ToolName.QES_CLASSIFICATION, {
    #     "message": "昨天的旅馆体验我非常满意".encode('unicode_escape'),
    #     "category_list": [
    #         {"category_id": "0", "category_desc": "技术问题".encode('unicode_escape')},
    #         {"category_id": "1", "category_desc": "娱乐问题".encode('unicode_escape')},
    #         {"category_id": "2", "category_desc": "商业问题".encode('unicode_escape')}
    #     ]
    # }).content[0].text
    # question_classification = eval(question_classification)
    # print(question_classification)
    rerank_result = exec_plugin(ToolName.RERANK, {
        "question": "香蕉是什么".encode('unicode_escape'),
        "references": [
            "香蕉是黄色的".encode('unicode_escape'),
            "香蕉是一种水果".encode('unicode_escape'),
            "今天天气很好".encode('unicode_escape'),
            "猴子爱吃香蕉".encode('unicode_escape')
        ]
    }).content[0].text
    rerank_result = eval(rerank_result)
    print(rerank_result["ranked_results"])
    for item in rerank_result["ranked_results"]:
        print(item["id"])
